################
#
# Deep Flow Prediction - N. Thuerey, K. Weissenov, H. Mehrotra, N. Mainali, L. Prantl, X. Hu (TUM)
#
# Compute errors for a test set and visualize. This script can loop over a range of models in 
# order to compute an averaged evaluation.
#
################

import os, sys, random, math
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import DataLoader

from dataset import TurbDataset
from DfpNet import TurbNetG, weights_init
import utils
from utils import log

suffix = ""  # 自定义加载和输出后缀
prefix = ""
if len(sys.argv) > 1:
    prefix = sys.argv[1]
    print("Output prefix: {}".format(prefix))

# 调整模型参数以匹配训练代码
expo = 7  # 与训练代码保持一致，增加模型容量
dataset = TurbDataset(None, mode=TurbDataset.TEST, dataDirTest="../data/test/")
testLoader = DataLoader(dataset, batch_size=1, shuffle=False)

# 调整输入张量以适应新增的管口掩码通道
targets = torch.FloatTensor(1, 3, 128, 128)
targets = Variable(targets)
targets = targets.cuda()
inputs = torch.FloatTensor(1, 4, 128, 128)  # 输入通道从3调整为4
inputs = Variable(inputs)
inputs = inputs.cuda()
masks = torch.FloatTensor(1, 1, 128, 128)  # 管口掩码
masks = Variable(masks)
masks = masks.cuda()

targets_dn = torch.FloatTensor(1, 3, 128, 128)
targets_dn = Variable(targets_dn)
targets_dn = targets_dn.cuda()
outputs_dn = torch.FloatTensor(1, 3, 128, 128)
outputs_dn = Variable(outputs_dn)
outputs_dn = outputs_dn.cuda()

# 初始化模型
netG = TurbNetG(channelExponent=expo)
lf = "./" + prefix + "testout{}.txt".format(suffix)
utils.makeDirs(["results_test"])

# 循环测试多个模型
avgLoss = 0.
losses = []
models = []

for si in range(25):
    s = chr(96 + si)
    if si == 0:
        s = ""  # 检查modelG和modelG + char
    modelFn = "./" + prefix + "modelG{}{}".format(suffix, s)
    if not os.path.isfile(modelFn):
        continue

    models.append(modelFn)
    log(lf, "Loading " + modelFn)
    netG.load_state_dict(torch.load(modelFn))
    log(lf, "Loaded " + modelFn)
    netG.cuda()

    criterionL1 = nn.L1Loss()
    criterionL1.cuda()
    L1val_accum = 0.0
    L1val_dn_accum = 0.0
    L1val_inlet_accum = 0.0  # 新增：管口区域的L1误差
    lossPer_p_accum = 0
    lossPer_v_accum = 0
    lossPer_accum = 0
    lossPer_p_inlet_accum = 0  # 新增：管口区域的压力相对误差

    netG.eval()

    for i, data in enumerate(testLoader, 0):
        inputs_cpu, targets_cpu, masks_cpu = data  # 接收管口掩码
        targets_cpu, inputs_cpu, masks_cpu = targets_cpu.float().cuda(), inputs_cpu.float().cuda(), masks_cpu.float().cuda()
        inputs.data.resize_as_(inputs_cpu).copy_(inputs_cpu)
        targets.data.resize_as_(targets_cpu).copy_(targets_cpu)
        masks.data.resize_as_(masks_cpu).copy_(masks_cpu)

        # 前向传播，传入管口掩码
        outputs = netG(inputs, masks)
        outputs_cpu = outputs.data.cpu().numpy()[0]
        targets_cpu_np = targets_cpu.cpu().numpy()[0]
        masks_cpu_np = masks_cpu.cpu().numpy()[0]

        # 计算整体L1误差（与训练一致，使用加权损失）
        pressure_weight = 0.7  # 与训练代码一致
        loss_p = criterionL1(outputs[:, 0] * masks, targets[:, 0] * masks)  # 管口加权
        loss_v = criterionL1(outputs[:, 1:], targets[:, 1:])
        lossL1 = pressure_weight * loss_p + (1 - pressure_weight) * loss_v
        L1val_accum += lossL1.item()

        # 单独计算管口区域的压力误差
        loss_p_inlet = criterionL1(outputs[:, 0] * masks, targets[:, 0] * masks)
        L1val_inlet_accum += loss_p_inlet.item()

        # 计算相对误差百分比
        lossPer_p = np.sum(np.abs(outputs_cpu[0] - targets_cpu_np[0])) / np.sum(np.abs(targets_cpu_np[0]))
        lossPer_v = (np.sum(np.abs(outputs_cpu[1] - targets_cpu_np[1])) + np.sum(np.abs(outputs_cpu[2] - targets_cpu_np[2]))) / (np.sum(np.abs(targets_cpu_np[1])) + np.sum(np.abs(targets_cpu_np[2])))
        lossPer = np.sum(np.abs(outputs_cpu - targets_cpu_np)) / np.sum(np.abs(targets_cpu_np))
        lossPer_p_accum += lossPer_p.item()
        lossPer_v_accum += lossPer_v.item()
        lossPer_accum += lossPer.item()

        # 计算管口区域的压力相对误差
        p_diff_inlet = np.abs(outputs_cpu[0] - targets_cpu_np[0]) * masks_cpu_np[0]
        p_target_inlet = np.abs(targets_cpu_np[0]) * masks_cpu_np[0]
        lossPer_p_inlet = np.sum(p_diff_inlet) / np.sum(p_target_inlet) if np.sum(p_target_inlet) != 0 else 0
        lossPer_p_inlet_accum += lossPer_p_inlet

        log(lf, "Test sample %d" % i)
        log(lf, "    pressure:  abs. difference, ratio: %f , %f " % (np.sum(np.abs(outputs_cpu[0] - targets_cpu_np[0])), lossPer_p.item()))
        log(lf, "    velocity:  abs. difference, ratio: %f , %f " % (np.sum(np.abs(outputs_cpu[1] - targets_cpu_np[1])) + np.sum(np.abs(outputs_cpu[2] - targets_cpu_np[2])), lossPer_v.item()))
        log(lf, "    aggregate: abs. difference, ratio: %f , %f " % (np.sum(np.abs(outputs_cpu - targets_cpu_np)), lossPer.item()))
        log(lf, "    pressure (inlet): abs. difference, ratio: %f , %f " % (np.sum(p_diff_inlet), lossPer_p_inlet))

        # 计算去归一化误差
        input_ndarray = inputs_cpu.cpu().numpy()[0]
        v_norm = (np.max(np.abs(input_ndarray[0, :, :]))**2 + np.max(np.abs(input_ndarray[1, :, :]))**2)**0.5

        outputs_denormalized = dataset.denormalize(outputs_cpu, v_norm)
        targets_denormalized = dataset.denormalize(targets_cpu_np, v_norm)

        outputs_denormalized_comp = np.array([outputs_denormalized])
        outputs_denormalized_comp = torch.from_numpy(outputs_denormalized_comp)
        targets_denormalized_comp = np.array([targets_denormalized])
        targets_denormalized_comp = torch.from_numpy(targets_denormalized_comp)

        targets_denormalized_comp, outputs_denormalized_comp = targets_denormalized_comp.float().cuda(), outputs_denormalized_comp.float().cuda()

        outputs_dn.data.resize_as_(outputs_denormalized_comp).copy_(outputs_denormalized_comp)
        targets_dn.data.resize_as_(targets_denormalized_comp).copy_(targets_denormalized_comp)

        lossL1_dn = criterionL1(outputs_dn, targets_dn)
        L1val_dn_accum += lossL1_dn.item()

        # 可视化输出，突出管口区域
        os.chdir("./results_test/")
        utils.imageOut("%04d" % (i), outputs_cpu, targets_cpu_np, mask=masks_cpu_np[0], normalize=False, saveMontage=True)  # 添加mask参数以突出管口
        os.chdir("../")

    log(lf, "\n")
    L1val_accum /= len(testLoader)
    lossPer_p_accum /= len(testLoader)
    lossPer_v_accum /= len(testLoader)
    lossPer_accum /= len(testLoader)
    L1val_dn_accum /= len(testLoader)
    L1val_inlet_accum /= len(testLoader)
    lossPer_p_inlet_accum /= len(testLoader)

    log(lf, "Loss percentage (p, v, combined): %f %%    %f %%    %f %% " % (lossPer_p_accum * 100, lossPer_v_accum * 100, lossPer_accum * 100))
    log(lf, "Loss percentage (pressure at inlet): %f %% " % (lossPer_p_inlet_accum * 100))
    log(lf, "L1 error: %f" % (L1val_accum))
    log(lf, "L1 error (pressure at inlet): %f" % (L1val_inlet_accum))
    log(lf, "Denormalized error: %f" % (L1val_dn_accum))
    log(lf, "\n")

    avgLoss += lossPer_accum
    losses.append(lossPer_accum)

if len(losses) > 1:
    avgLoss /= len(losses)
    lossStdErr = np.std(losses) / math.sqrt(len(losses))
    log(lf, "Averaged relative error and std dev across models:   %f , %f " % (avgLoss, lossStdErr))